import copy
import numpy as np
import matplotlib.pyplot as plt
from turtle import *


from double_s import DoubleS
from path import PathButterfly


class KinematicScaling:
    def __init__(self, double_s: DoubleS, total_time: float):
        self.double_s = copy.deepcopy(double_s)
        self.total_time = total_time
        self.factor = self.double_s.get_total_time() / total_time

    def get_trajectory(self, t):
        trajectory = self.double_s.get_trajectory(t * self.factor)
        return (trajectory[0],
                trajectory[1] * self.factor,
                trajectory[2] * self.factor ** 2,
                trajectory[3] * self.factor ** 3)

    def get_total_time(self):
        return self.total_time


if __name__ == '__main__':
    q0 = np.array(0.0)
    q1 = np.array(np.pi * 2)
    vmax = np.array(5.0)
    amax = np.array(10.0)
    jmax = np.array(30.0)

    num = 1001

    double_s = DoubleS(q0, q1, vmax, amax, jmax)
    kinematic_scaling = KinematicScaling(double_s, 5)
    path_butterfly = PathButterfly(60)

    print('double S total time: ', double_s.get_total_time())
    print('kinematic scaling total time: ', kinematic_scaling.get_total_time())

    q_double_s = []
    dq_double_s = []
    ddq_double_s = []
    dddq_double_s = []
    q_kinematic_scaling = []
    dq_kinematic_scaling = []
    ddq_kinematic_scaling = []
    dddq_kinematic_scaling = []
    t = np.linspace(0, np.ceil(np.max((double_s.get_total_time(), kinematic_scaling.get_total_time()))), num)

    for ti in t:
        trajectory_double_s = double_s.get_trajectory(ti)
        q_double_s.append(trajectory_double_s[0])
        dq_double_s.append(trajectory_double_s[1])
        ddq_double_s.append(trajectory_double_s[2])
        dddq_double_s.append(trajectory_double_s[3])

        trajectory_kinematic_scaling = kinematic_scaling.get_trajectory(ti)
        q_kinematic_scaling.append(trajectory_kinematic_scaling[0])
        dq_kinematic_scaling.append(trajectory_kinematic_scaling[1])
        ddq_kinematic_scaling.append(trajectory_kinematic_scaling[2])
        dddq_kinematic_scaling.append(trajectory_kinematic_scaling[3])

    plt.figure(1)

    plt.subplot(2, 2, 1)
    plt.plot(t, q_double_s, label='double S')
    plt.plot(t, q_kinematic_scaling, label='kinematic scaling')
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$q (rad)$')
    plt.legend()

    plt.subplot(2, 2, 2)
    plt.plot(t, dq_double_s, label='double S')
    plt.plot(t, dq_kinematic_scaling, label='kinematic scaling')
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$\dot{q} (rad/s)$')
    plt.legend()

    plt.subplot(2, 2, 3)
    plt.plot(t, ddq_double_s, label='double S')
    plt.plot(t, ddq_kinematic_scaling, label='kinematic scaling')
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$\ddot{q} (rad/s^2)$')
    plt.legend()

    plt.subplot(2, 2, 4)
    plt.plot(t, dddq_double_s, label='double S')
    plt.plot(t, dddq_kinematic_scaling, label='kinematic scaling')
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$\dddot{q} (rad/s^3)$')
    plt.legend()

    plt.tight_layout()
    plt.show()

    pensize(3)
    pencolor("blue")
    for qi in q_kinematic_scaling:
        goto(*(path_butterfly.get_path(qi)[0]))
